import rclpy
from rclpy.node import Node
from tf2_ros import TransformListener, Buffer
import transforms3d as tfs
import numpy as np
from xarm.wrapper import XArmAPI # type: ignore
from scipy.spatial.transform import Rotation as R


'''
1.通过控制器上的Aruco tag的xyz差值,控制机械臂末端的x,y,z三个自由度
  通过控制器上的Aruco tag的rpy差值,控制机械臂末端的r,p,y三个自由度
  上述两者可以实现对整个机械臂末端控制器的运动控制
2.使用此代码要修改的地方:
    1) ip = '192.168.1.231'
    2) 初始位置 initial_position_ = np.array([477, 1.4, 416, 180, 0.6, 90])
'''



# ip 与特定机械臂有关
ip = '192.168.1.228'

# 实例化机械臂
arm = XArmAPI(ip)
 
# 初始化机械臂
def init_arm():
    arm.clean_error()
    arm.motion_enable(enable=True)
    arm.set_mode(7)
    arm.set_state(state=0)    
    
class TransformListenerNode(Node):
    def __init__(self):
        super().__init__('transform_listener_node')

        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)

        # 你可以设置一个定时器来定期检查变换是否可用
        self.timer = self.create_timer(delta_time_, self.on_timer)

    def on_timer(self):
        # 检查变换是否可用
        flag_once =  True
        try:
            # 这里你需要替换成你的实际坐标系名称
            source_frame = 'tag'    # tag末端坐标系
            target_frame = 'fake_origin'   #虚拟零点，按下小乌龟上键后固定次坐标系

            # 获取当前时间的变换
            transform = self.tf_buffer.lookup_transform(target_frame, source_frame, rclpy.time.Time())
            #获取转换数据
            translation = transform.transform.translation
            rotation = transform.transform.rotation
            # 将旋转转换为四元数形式的NumPy数组
            rotation_array = np.array([rotation.w, rotation.x, rotation.y, rotation.z])
            # rotation_bias_array = np.array([bias_rotation.w, bias_rotation.x, bias_rotation.y, bias_rotation.z])
            # 四元数转固定轴欧拉角
            euler_angles_from_quat = tfs.euler.quat2euler(rotation_array, "sxyz")   
            # start calculate
            pose = np.array([translation.z, -translation.x, -translation.y,
                             -euler_angles_from_quat[0], -euler_angles_from_quat[2], -euler_angles_from_quat[1]])
            # 设定死区
            if -0.05 < pose[0] < 0.05:
                pose[0] = 0.0
            if -0.05 < pose[1] < 0.05:
                pose[1] = 0.0   
            if -0.05 < pose[2] < 0.05:
                pose[2] = 0.0
            if -0.01 < pose[3] < 0.01:
                pose[3] = 0.0
            if -0.01 < pose[4] < 0.01:
                pose[4] = 0.0
            if -0.01 < pose[5] < 0.01:    
                pose[5] = 0.0
            # mapping
            pose = pose * 100
            pose[0] *= 6
            pose[1] *= 6
            pose[2] *= 6
            #calculate the final position
            final_position_ = initial_position_ + pose
            # 设定
            arm.set_position(final_position_[0], final_position_[1], final_position_[2], 
                             final_position_[3], final_position_[4], final_position_[5], 
                             speed=10000000, mvacc=1000,cmvtime=delta_time_)   
        except Exception as e:
            self.get_logger().warn(f'Could not transform: {e}')

def main(args=None):
    init_arm()
    rclpy.init(args=args)
    node = TransformListenerNode()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        rclpy.shutdown()

if __name__ == '__main__':
    # delta位姿
    pose = np.array([0,0,0,0,0,0])
    # 频率
    delta_time_ = 0.01
    # 初始位置 与特定机械臂有关
    initial_position_ = np.array([477, 1.4, 416, 180, 0.6, 90])
    main()